# -*- coding: utf-8 -*-
"""
GIS工具模块

此模块提供用于处理GIS数据和省份边界信息的函数。
主要功能包括：
- 读取省份边界GIS文件
- 提取省份边界数据
- 判断点是否在省份内
"""

import os
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point


def read_province_shapefile(shp_path):
    """
    读取省份边界GIS文件
    
    参数:
        shp_path (str): 省份边界GIS文件路径
        
    返回:
        geopandas.GeoDataFrame: 包含省份边界数据的GeoDataFrame
    """
    try:
        # 读取shapefile文件
        gdf = gpd.read_file(shp_path, encoding='utf-8')
        print(f"成功读取省份边界GIS文件，共{len(gdf)}个省份")
        return gdf
    except Exception as e:
        print(f"读取省份边界GIS文件失败: {e}")
        return None


def filter_provinces(province_gdf, provinces, name_field='NAME'):
    """
    筛选指定的省份
    
    参数:
        province_gdf (geopandas.GeoDataFrame): 包含省份边界数据的GeoDataFrame
        provinces (list): 要筛选的省份名称列表
        name_field (str): 省份名称字段，默认为'NAME'
        
    返回:
        geopandas.GeoDataFrame: 筛选后的GeoDataFrame
    """
    try:
        # 筛选指定的省份
        filtered_gdf = province_gdf[province_gdf[name_field].isin(provinces)]
        print(f"成功筛选{len(filtered_gdf)}个省份")
        return filtered_gdf
    except Exception as e:
        print(f"筛选省份失败: {e}")
        return None


def get_province_bounds(province_gdf, province, name_field='NAME'):
    """
    获取指定省份的边界范围
    
    参数:
        province_gdf (geopandas.GeoDataFrame): 包含省份边界数据的GeoDataFrame
        province (str): 省份名称
        name_field (str): 省份名称字段，默认为'NAME'
        
    返回:
        tuple: 省份边界范围(min_lon, min_lat, max_lon, max_lat)
    """
    try:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 获取省份边界范围
        bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
        min_lon, min_lat, max_lon, max_lat = bounds
        
        print(f"{province}省份边界范围: 经度({min_lon:.4f}, {max_lon:.4f}), 纬度({min_lat:.4f}, {max_lat:.4f})")
        return min_lon, min_lat, max_lon, max_lat
    except Exception as e:
        print(f"获取{province}省份边界范围失败: {e}")
        return None


def is_point_in_province(province_gdf, province, lon, lat, name_field='NAME'):
    """
    判断点是否在指定省份内
    
    参数:
        province_gdf (geopandas.GeoDataFrame): 包含省份边界数据的GeoDataFrame
        province (str): 省份名称
        lon (float): 经度
        lat (float): 纬度
        name_field (str): 省份名称字段，默认为'NAME'
        
    返回:
        bool: 点是否在省份内
    """
    try:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 创建点对象
        point = Point(lon, lat)
        
        # 判断点是否在省份内
        return province_geom.contains(point)
    except Exception as e:
        print(f"判断点是否在{province}省份内失败: {e}")
        return False


def extract_province_grid_points(province_gdf, lons, lats, name_field='NAME', provinces=None):
    """
    提取指定省份的观测点经纬度数据
    
    参数:
        province_gdf (geopandas.GeoDataFrame): 省份边界GeoDataFrame
        lons (numpy.ndarray): 经度数组
        lats (numpy.ndarray): 纬度数组
        name_field (str): 省份名称字段，默认为'NAME'
        provinces (list): 要提取的省份名称列表，如为None则提取所有省份
        
    返回:
        dict: 包含各省份观测点经纬度数据的字典
    """
    import numpy as np
    
    # 如果未指定省份，则使用所有省份
    if provinces is None:
        provinces = province_gdf[name_field].tolist()
    else:
        # 筛选指定的省份
        filtered_gdf = filter_provinces(province_gdf, provinces, name_field)
        provinces = filtered_gdf[name_field].tolist()
    
    if not provinces:
        print("没有有效的省份可以提取")
        return {}
    
    # 创建结果字典
    result = {}
    
    # 对每个省份提取观测点
    for province in provinces:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 获取省份边界范围
        bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
        min_lon, min_lat, max_lon, max_lat = bounds
        
        # 提取在省份范围内的经纬度点
        province_lons = lons[(lons >= min_lon) & (lons <= max_lon)]
        province_lats = lats[(lats >= min_lat) & (lats <= max_lat)]
        
        # 创建经纬度网格点
        lon_grid, lat_grid = np.meshgrid(province_lons, province_lats)
        
        # 将网格点展平为一维数组
        grid_points_flat = np.column_stack((lon_grid.flatten(), lat_grid.flatten()))
        
        # 筛选真正在省份多边形内的点
        in_province_points = []
        for point in grid_points_flat:
            if province_geom.contains(Point(point[0], point[1])):
                in_province_points.append(point)
        
        # 转换为numpy数组
        if in_province_points:
            grid_points = np.array(in_province_points)
        else:
            grid_points = np.empty((0, 2))
        
        # 保存到结果字典
        result[province] = {
            'grid_points': grid_points,
            'lon_count': len(province_lons),
            'lat_count': len(province_lats),
            'total_points': len(grid_points)
        }
        
        print(f"省份: {province}, 经度点数: {len(province_lons)}, 纬度点数: {len(province_lats)}, 有效点数: {len(grid_points)}")
    
    return result